import httpx
import os
import pandas as pd
import gradio as gr
import plotly.graph_objects as go

from xhtml2pdf import pisa
from jinja2 import Environment, FileSystemLoader


from volcenginesdkarkruntime import Ark
from typing import Dict, List, Any


client = Ark(
    api_key=os.environ.get("ARK_API_KEY"),
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    timeout=httpx.Timeout(timeout=1800),
)

try:
    # 没有去wps云盘下载
    pth = r"..\world_data\OECD.SDD.NAD,DSD_NASEC20@DF_T720R_A,+all.csv"
    data = pd.read_csv(pth, encoding='gbk')
except:
    pth=r"D:\WPS云文件储存位置\252320705\WPS云盘\云数据\OECD国家资产负债表\OECD.SDD.NAD,DSD_NASEC20@DF_T720R_A,+all.csv"
    # pth = r"C:\Users\59567\OneDrive\Desktop\OECD国家资产负债表\OECD.SDD.NAD,DSD_NASEC20@DF_T720R_A,+all.csv"
    pth = r"../../云数据/OECD国家资产负债表/OECD.SDD.NAD,DSD_NASEC20@DF_T720R_A,+all.csv"
    data = pd.read_csv(pth, encoding='gbk')


auto_choice = [('UNIT_MEASURE', 'USD')]
for name, value in auto_choice:
    data = data[data[name] == value]

assert isinstance(data, pd.DataFrame)

filter_data_raw = data.loc[:,['REF_AREA', 'TIME_PERIOD', 'Institutional sector', 
                             'Financial instruments and non-financial assets', 'Original and residual maturity', 
                             'ACCOUNTING_ENTRY',
                             'OBS_VALUE']]

user_choice = [
    'REF_AREA',
    'Institutional sector', 
    'Financial instruments and non-financial assets',
    'Original and residual maturity',
    'ACCOUNTING_ENTRY',
    'TIME_PERIOD',
    ]

choice = {name: data[name].unique() for name in user_choice}

print(data.shape)
for k,v in choice.items():
    print(k, v)

country_en_to_cn ={
    'AUT': '奥地利',
    'BEL': '比利时',
    'CAN': '加拿大',
    'CHL': '智利',
    'CZE': '捷克',
    'DNK': '丹麦',
    'EST': '爱沙尼亚',
    'FIN': '芬兰',
    'FRA': '法国',
    'DEU': '德国',
    'GRC': '希腊',
    'HUN': '匈牙利',
    'ISL': '冰岛',
    'IRL': '爱尔兰',
    'ISR': '以色列',
    'ITA': '意大利',
    'JPN': '日本',
    'KOR': '韩国',
    'LUX': '卢森堡',
    'MEX': '墨西哥',
    'NLD': '荷兰',
    'NZL': '新西兰',
    'NOR': '挪威',
    'POL': '波兰',
    'PRT': '葡萄牙',
    'ESP': '西班牙',
    'SWE': '瑞典',
    'CHE': '瑞士',
    'TUR': '土耳其',
    'GBR': '英国',
    'USA': '美国',
    }
    

institution_en_to_cn = {
    'Financial corporations': '金融机构',
    'Households and non-profit institutions serving households (NPISH)': '家庭和非盈利机构',
    'General government': '广义政府',
    'Central bank': '中央银行',
    'Total economy': '整个经济',
    'Non-financial corporations': '非金融机构',
    'Rest of the world': '世界其他',
}

financial_en_to_cn = {
    'Total financial instruments': '总金融项目',
    'Monetary gold and Special Drawing Rights (SDRs)': '货币黄金和特别提款权（SDRs）',
    'Currency and deposits': '货币和存款',
    'Debt securities': '债务证券',
    'Loans': '贷款',
    'Equity and investment fund shares/units': '股权和投资基金',
    'Insurance, pension and standardized guarantee schemes': '保险、养老金和标准化担保计划',
    'Financial derivatives and employee stock options': '金融衍生品和员工股票期权',
    'Other accounts receivable/payable': '其他应收/应付账款',
    'Net financial worth': '净金融资产',
}


def convert_html_to_pdf(source_html, output_filename):
    # 打开一个文件以写入 PDF 内容
    result_file = open(output_filename, "w+b")
    # 使用 pisa 创建 PDF
    pisa_status = pisa.CreatePDF(
        source_html,
        dest=result_file,
        encoding='utf-8')
    # 关闭文件
    result_file.close()

def df_to_pdf(df):
    # 配置Jinja2环境
    env = Environment(loader=FileSystemLoader('.'))
    template = env.get_template('./生成报告/report_template.html')

    # 准备数据以填充模板
    columns = df.columns.tolist()
    rows = df.values.tolist()

    # 渲染模板
    output = template.render(columns=columns, rows=rows)

    # 保存为HTML文件
    output_file_path = 'automated_report.html'
    with open(output_file_path, 'w', encoding='utf-8') as f:
        f.write(output)

    print(f"自动化报告已保存为 {output_file_path}")
  
    convert_html_to_pdf(output, 'automated_report.pdf')

class ShowData:
    def __init__(self):
        self.country_cn_to_en = {v:k for k,v in country_en_to_cn.items()}
        self.country_cn = list(self.country_cn_to_en.keys())
        # Bug修复：添加类型转换
        self.countrys = pd.Series(data['REF_AREA']).unique().tolist()
        self.institution_cn_to_en = {v:k for k,v in institution_en_to_cn.items()}
        self.institutions = [v for k,v in institution_en_to_cn.items()]
        self.years_start = list(range(1995,2024))
        self.years_end = list(range(1995,2024))
        self.financial_cn_to_en = {v:k for k,v in financial_en_to_cn.items()}
        self.financial = [v for k,v in financial_en_to_cn.items()]
        # Bug修复：添加类型转换
        self.maturity = pd.Series(data['Original and residual maturity']).unique().tolist()
        # Bug修复：添加类型转换
        self.account = pd.Series(data['ACCOUNTING_ENTRY']).unique().tolist()
        self.plot_data = pd.DataFrame()
    
    def save_fig(self, save_path):
        # 保存为 HTML 文件内容
        html_content = self.fig.to_html(full_html=False, include_plotlyjs='cdn')
        # 添加字符编码元标签
        html_content = '<!DOCTYPE html><html><head><meta charset="UTF-8"></head><body>' + html_content + '</body></html>'
        # 将内容写入文件
        with open(f'{save_path}', 'w', encoding='utf-8') as f:
            f.write(html_content)
        print(f"save done! {save_path}")
        return gr.Info(f"保存{save_path}成功！")
        
    def plot_data_fn(self, idx_, col_):
        data = self.plot_data
        fig = go.Figure()
        pivoted_df = data.pivot_table(index=idx_, columns=col_, values='value')
        try:
            pivoted_df.to_excel('../pivot.xlsx')
        except:
            pass
        pivoted_df.fillna(0, inplace=True)
        for col in pivoted_df.columns:
            fig.add_trace(go.Bar(x=pivoted_df.index.to_list(), y=pivoted_df[col].to_list(), name=col))
        fig.update_layout(
            barmode='group',
            height=600,
            legend=dict(
                orientation="h",  # 水平方向排列
                yanchor="bottom",
                y=1.02,  # 适当调整y值，使图例在图表上方
                xanchor="center",
                x=0.5  # 使图例居中
            )
        )
        self.fig = fig   
        return fig
        
    def filter_data(self, country, year_start, year, institution, financial, maturity, account):
        # data = data[data['REF_AREA'] == country]
        filter_df = filter_data_raw.copy()
        
        country_en = [self.country_cn_to_en[k] for k in country]
        institution_en = [self.institution_cn_to_en[k] for k in institution]
        financial_en = [self.financial_cn_to_en[k] for k in financial]
        filter_dic = {'REF_AREA': country_en, 
                      'TIME_PERIOD': list(range(year_start, year+1)), 
                      'Institutional sector': institution_en, 
                      'Financial instruments and non-financial assets': financial_en, 
                      'Original and residual maturity': maturity, 
                      'ACCOUNTING_ENTRY': account, 
                      }
        filter_one = {k: v for k, v in filter_dic.items() if len(v) == 1}
        filter_more = {k:v for k,v in filter_dic.items() if len(v) > 1}
        
        filter_text_ = ''
        for k,v in filter_one.items():
            print('filter', k, v[0], filter_df.shape)
            filter_df = filter_df[filter_df[k] == v[0]]
            filter_text_ += f'{k}:{v[0]};筛选后还剩下{filter_df.shape[0]}条数据\n'
            print('after_filter', k, v[0], filter_df.shape)
        for k,v in filter_more.items():
            print('filter', k, v, filter_df.shape)
            filter_df = filter_df[filter_df[k].isin(v)]
            filter_text_ += f'{k}:{v};筛选后还剩下{filter_df.shape[0]}条数据\n'
            filter_text_ += f'单位：万亿美元\n'
            print('after_filter', k, v, filter_df.shape)
        
        show_column = [s for s,_ in filter_more.items()]
        show_column.append('OBS_VALUE')
        filter_df = filter_df.loc[:, show_column]
        filter_df['value'] = filter_df['OBS_VALUE'] / 1000000
        filter_df['value'] = filter_df['value'].round(2)
        if 'TIME_PERIOD' in filter_df.columns:
            filter_df['TIME_PERIOD'] = filter_df['TIME_PERIOD'].astype(int)
        if 'OBS_VALUE' in filter_df.columns:
            del filter_df['OBS_VALUE']
        if 'REF_AREA' in filter_df.columns:
            filter_df['国家'] = filter_df['REF_AREA'].map(country_en_to_cn)
            del filter_df['REF_AREA']
        self.plot_data = filter_df
        cols = filter_df.columns.to_list()
        return filter_df, gr.Text(
            value=filter_df.shape[0], label='数据量'), gr.Text(
            value=filter_text_, label='筛选过程'), gr.Radio(choices=cols,  value=cols[-1]), gr.Radio(choices=cols,  value=cols[0])

    def ana_fn(self, hint, model_select):
        data = self.plot_data
        df_text = data.to_csv(sep='\t', na_rep='')
        prompt = f"""现在有这些资产负债表数据：数据说明如下{hint}\n{df_text}\n请给出分析结果"""
        stream = client.chat.completions.create(
            model=model_select,
            messages=[
                {"role": "system", "content": "你是懂国家金融资产负债表的专业人工智能助手"},
                {"role": "user", "content": prompt},
            ],
            stream=True
        )
        result = ""
        for chunk in stream:
            # Bug 修复：检查 chunk 是否为具有 choices 属性的对象
            if not chunk.choices:
                continue
            content = chunk.choices[0].delta.content
            if content is not None:
                result += content
            yield result

    def country_select_all_fn(self):
        return gr.CheckboxGroup(choices=self.country_cn, value=self.country_cn)

    def country_deselect_all_fn(self):
        return gr.CheckboxGroup(choices=self.country_cn, value=[])
    
    def institution_select_all_fn(self):
        return gr.CheckboxGroup(choices=self.institutions, value=self.institutions)
    
    def institution_deselect_all_fn(self):
        return gr.CheckboxGroup(choices=self.institutions, value=[])
    
    def financial_select_all_fn(self):
        return gr.CheckboxGroup(choices=self.financial, value=self.financial)

    def financial_deselect_all_fn(self):
        return gr.CheckboxGroup(choices=self.financial, value=[])
    
    def html(self):
        with gr.Blocks() as demo:
            gr.Markdown("# 多国金融资产负债表对比图")
            with gr.Tab("选择数据"):
                gr.Markdown("## 选择国家")
                countrys = gr.CheckboxGroup(choices=self.country_cn, label="国家", value=['美国'], interactive=True)
                with gr.Row():
                    country_select_all = gr.Button("全选", scale=2)
                    country_deselect_all = gr.Button("全不选", scale=2)
                country_select_all.click(self.country_select_all_fn, inputs=[], outputs=countrys)
                country_deselect_all.click(self.country_deselect_all_fn, inputs=[], outputs=countrys)
                    
                gr.Markdown("## 选择年份")
                with gr.Row():
                    years_start = gr.Slider(minimum=1995, maximum=2023, step=1, label="开始年份", value=2000, interactive=True)
                    years = gr.Slider(minimum=1995, maximum=2023, step=1, label="年份", value=2022, interactive=True)
                gr.Markdown("## 金融机构与金融资产")
                with gr.Row():
                    institutions = gr.CheckboxGroup(choices=self.institutions, label="金融机构", value=['整个经济'], interactive=True, scale=2)
                    with gr.Column():
                        institutions_select_all = gr.Button("全选", scale=2)
                        institutions_deselect_all = gr.Button("全不选", scale=2)
                        institutions_select_all.click(self.institution_select_all_fn, inputs=[], outputs=institutions)
                        institutions_deselect_all.click(self.institution_deselect_all_fn, inputs=[], outputs=institutions)
                with gr.Row():        
                    financial = gr.CheckboxGroup(choices=self.financial, label="金融资产", value=['总金融项目'], interactive=True, scale=3)
                    with gr.Column():
                        financial_select_all = gr.Button("全选", scale=2)
                        financial_deselect_all = gr.Button("全不选", scale=2)
                        financial_select_all.click(self.financial_select_all_fn, inputs=[], outputs=financial)
                        financial_deselect_all.click(self.financial_deselect_all_fn, inputs=[], outputs=financial)
                with gr.Row():
                    text_return = gr.Textbox(label="符合条件的数据", interactive=False, scale=3)
                    bt = gr.Button("我已选好，筛选数据")
                gr.Markdown("## 选择期限与会计科目")
                with gr.Row():
                    maturity = gr.CheckboxGroup(choices=self.maturity, label="期限", value='Not applicable', interactive=True, scale=3)
                    account = gr.CheckboxGroup(choices=self.account, label="会计科目", value='A', interactive=True, scale=1)
            
                gr.Markdown("## 选择图表")
                
                with gr.Row():
                    x_axis_select = gr.Radio(label='X轴设定', interactive=True, scale=2)
                    second_x_axis_select = gr.Radio(label='第二X轴设定', interactive=True, scale=2)
                    bt_plot = gr.Button("画图", scale=1)
                    
                xy_img = gr.Plot(label="图表")
                
                with gr.Row():
                    save_path = gr.Textbox(label="保存路径", interactive=True, value="C:\\Users\\59567\\save_plot.html", scale=3)
                    bt_report = gr.Button("保存此图", scale=1)
            
            with gr.Tab("展示数据"):
                gr.Markdown("## 筛选条件")
                filter_btn = gr.Button("筛选数据")
                filter_text = gr.Textbox(label="筛选条件", interactive=False, lines=6)
                show = gr.DataFrame()
                gr.Markdown("## 基于大预言模型，分析数据")
                with gr.Row():
                    input_text = gr.Textbox(label="请输入辅助说明", lines=1, interactive=True)
                    ana_bt = gr.Button("尝试分析")
                model_choice = gr.Dropdown(
                    choices=["doubao-1-5-pro-256k-250115", "deepseek-r1-250120"], 
                    label="选择模型", value="doubao-1-5-pro-256k-250115")
                llm_text = gr.Markdown(label="分析结果")
                
            
            ana_bt.click(self.ana_fn, inputs=[input_text, model_choice], outputs=[llm_text])
                
            gr.on(
                triggers=[bt_report.click], inputs=[save_path], fn=self.save_fig)

            gr.on(
                triggers=[bt_plot.click], inputs=[x_axis_select, second_x_axis_select],
                outputs=[xy_img], fn=self.plot_data_fn)

            gr.on(
                triggers=[bt.click, filter_btn.click], 
                  inputs=[countrys, years_start, years, institutions, financial, maturity, account], 
                  outputs=[show, text_return, filter_text, x_axis_select, second_x_axis_select], 
                  fn=self.filter_data)
        
        return demo
    
sd = ShowData()
sd.html().launch()